package main

import (
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/creack/pty"
	"github.com/gorilla/websocket"
	"github.com/isayme/go-logger"
	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/leaanthony/mewn"
	"github.com/vmihailenco/msgpack/v5"
)

const (
	TYPE_DATA   = 0
	TYPE_RESIZE = 1
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type windowSize struct {
	Rows uint16 `json:"rows" msgpack:"rows"`
	Cols uint16 `json:"cols" msgpack:"cols"`
	X    uint16 `json:"x" msgpack:"x"`
	Y    uint16 `json:"y" msgpack:"y"`
}

var contextPath string
var port int
var ttyCmd string

// www
var indexHtml string
var uploadHtml string
var indexJs string
var indexCss string
var faviconSvg string

func main() {
	flag.StringVar(&contextPath, "c", "", "contextPath")
	flag.IntVar(&port, "p", 10081, "port")
	flag.StringVar(&ttyCmd, "s", "bash", "shell")
	flag.Parse()

	e := echo.New()

	e.Use(middleware.Logger())
	e.Use(middleware.Recover())
	e.GET(contextPath+"/ws", handleWebsocket)
	indexHtml = mewn.String("./www/index.html")
	indexHtml = strings.ReplaceAll(indexHtml, "{contextPath}", contextPath)
	uploadHtml = mewn.String("./www/upload.html")
	uploadHtml = strings.ReplaceAll(uploadHtml, "{contextPath}", contextPath)
	indexJs = mewn.String("./www/index.js")
	indexCss = mewn.String("./www/index.css")
	faviconSvg = mewn.String("./www/favicon.svg")

	handleReq(e)

	e.Logger.Fatal(e.Start(":" + strconv.Itoa(port)))
}

func handleReq(e *echo.Echo) {
	e.GET(contextPath, func(c echo.Context) error {
		c.Response().Header().Set(echo.HeaderContentType, echo.MIMETextHTMLCharsetUTF8)
		return c.String(http.StatusOK, indexHtml)
	})
	e.GET(contextPath+"/", func(c echo.Context) error {
		c.Response().Header().Set(echo.HeaderContentType, echo.MIMETextHTMLCharsetUTF8)
		return c.String(http.StatusOK, indexHtml)
	})
	e.GET(contextPath+"/index.js", func(c echo.Context) error {
		c.Response().Header().Set(echo.HeaderContentType, echo.MIMEApplicationJavaScriptCharsetUTF8)
		return c.String(http.StatusOK, indexJs)
	})
	e.GET(contextPath+"/index.css", func(c echo.Context) error {
		c.Response().Header().Set(echo.HeaderContentType, "text/css; charset=utf-8")
		return c.String(http.StatusOK, indexCss)
	})
	e.GET(contextPath+"/favicon.svg", func(c echo.Context) error {
		c.Response().Header().Set(echo.HeaderContentType, "image/svg+xml")
		return c.String(http.StatusOK, faviconSvg)
	})
	e.GET(contextPath+"/upload", func(c echo.Context) error {
		c.Response().Header().Set(echo.HeaderContentType, echo.MIMETextHTMLCharsetUTF8)
		return c.String(http.StatusOK, uploadHtml)
	})
	e.POST(contextPath+"/upload", func(c echo.Context) error {
		// 从表单中获取文件
		file, err := c.FormFile("file")
		if err != nil {
			return c.JSON(http.StatusBadRequest, map[string]string{
				"error": "未找到上传的文件",
			})
		}

		// 获取临时目录
		tmpDir := os.TempDir()
		if tmpDir == "" {
			return c.JSON(http.StatusInternalServerError, map[string]string{
				"error": "无法获取临时目录",
			})
		}

		// 构建目标文件路径
		dstPath := filepath.Join(tmpDir, file.Filename)

		// 检查文件是否已存在，如果存在则删除
		if _, err := os.Stat(dstPath); err == nil {
			if err := os.Remove(dstPath); err != nil {
				return c.JSON(http.StatusInternalServerError, map[string]string{
					"error": fmt.Sprintf("无法删除已存在的文件: %v", err),
				})
			}
		}

		// 打开上传的文件
		src, err := file.Open()
		if err != nil {
			return c.JSON(http.StatusInternalServerError, map[string]string{
				"error": fmt.Sprintf("无法打开上传的文件: %v", err),
			})
		}
		defer src.Close()

		// 创建目标文件
		dst, err := os.Create(dstPath)
		if err != nil {
			return c.JSON(http.StatusInternalServerError, map[string]string{
				"error": fmt.Sprintf("无法创建目标文件: %v", err),
			})
		}
		defer dst.Close()

		// 复制文件内容
		if _, err = io.Copy(dst, src); err != nil {
			return c.JSON(http.StatusInternalServerError, map[string]string{
				"error": fmt.Sprintf("无法复制文件内容: %v", err),
			})
		}

		// 返回成功响应
		return c.JSON(http.StatusOK, map[string]string{
			"message": "文件上传成功",
			"path":    dstPath,
		})
	})
}

func handleWebsocket(c echo.Context) error {
	ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		return err
	}
	defer ws.Close()

	shell := exec.Command("sh", "-c", ttyCmd)
	shell.Env = append(shell.Environ(), "TERM=xterm")

	ptmx, err := pty.Start(shell)
	if err != nil {
		logger.Errorf("pty.Start error: %v", err)
		return err
	}

	defer func() { ptmx.Close() }()

	wg := sync.WaitGroup{}
	wg.Add(2)

	go func() {
		defer wg.Done()

		dataTypeBuf := make([]byte, 1)

		for {
			messageType, reader, err := ws.NextReader()
			if err != nil {
				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					logger.Errorf("unable to grab next reader: %v", err)
				}
				return
			}

			if messageType == websocket.TextMessage {
				warnMsg := fmt.Sprintf("unexpected message type: %d", messageType)
				logger.Warn(warnMsg)
				return
			}

			read, err := reader.Read(dataTypeBuf)
			if err != nil {
				errorMsg := fmt.Sprintf("unable to read message type from reader: %v", err)
				logger.Error(errorMsg)
				return
			}

			if read != 1 {
				logger.Error("read data type fail")
				return
			}

			dataType := dataTypeBuf[0]
			switch dataType {
			case TYPE_DATA:
				copied, err := io.Copy(ptmx, reader)
				if err != nil {
					logger.Errorf("error after copying %d bytes, err: %v", copied, err)
					return
				}
			case TYPE_RESIZE:
				resizeMsgBuf, err := io.ReadAll(reader)
				if err != nil {
					logger.Warnf("error decoding resize message: %v", err)
					continue
				}

				resizeMessage := windowSize{}
				err = msgpack.Unmarshal(resizeMsgBuf, &resizeMessage)
				if err != nil {
					logger.Warnf("error msgpack.Unmarshal: %v", err)
					continue
				}

				logger.Infof("resizing terminal: %+v", resizeMessage)

				winSize := &pty.Winsize{
					Rows: resizeMessage.Rows,
					Cols: resizeMessage.Cols,
					X:    resizeMessage.X,
					Y:    resizeMessage.Y,
				}
				pty.Setsize(ptmx, winSize)
			default:
				logger.Errorf("unknown data type: %d", dataTypeBuf[0])
			}
		}
	}()

	go func() {
		defer wg.Done()

		readBuf := make([]byte, 1024)

		for {
			n, err := ptmx.Read(readBuf)
			if n >= 0 {
				err := ws.WriteMessage(websocket.BinaryMessage, readBuf[:n])
				if err != nil {
					logger.Errorf("receive error: %v", err)
					break
				}
			}

			if err != nil {
				if err != io.EOF {
					logger.Errorf("read ptmx error: %v", err)
				}
				break
			}
		}
	}()

	wg.Wait()

	return nil
}
